package cc.shimeng.fucloud.controller.ai;

import cc.shimeng.fucloud.config.LoggingAdvisor;
import cc.shimeng.fucloud.tools.ClassesTool;
import cc.shimeng.fucloud.tools.GradeTool;
import cc.shimeng.fucloud.tools.StudentTool;
import cc.shimeng.fucloud.tools.SubjectTool;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.LocalDate;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

@CrossOrigin
@RestController
public class AiController {

    //    @Autowired
//    private VectorStore vectorStore;
    private final ChatClient chatClient;

    public AiController(ChatClient.Builder chatClientBuilder,
//                        VectorStore vectorStore,
                        ChatMemory chatMemory,
                        ToolCallbackProvider mcpTools,
                        ClassesTool classesTool,
                        SubjectTool subjectTool,
                        GradeTool gradeTool,
                        StudentTool studentTool
    ) {
        this.chatClient = chatClientBuilder
                .defaultSystem("""
                            您是“三才学院”的客户聊天支持代理。请以友好、乐于助人且愉快的方式来回复。
                                          您正在通过在线聊天系统与客户互动。
                                          请讲中文。
                                          今天的日期是 {current_date}.
                                           在新增、修改、删除“成绩grade”function-call前，请先获取学生信息和科目信息，并且一定要等用户回复"确定"之后才进行新增、修改、删除的function-call。 
                        """)
                .defaultAdvisors(
                        new PromptChatMemoryAdvisor(chatMemory),
                        new LoggingAdvisor()
                )
                .defaultTools(mcpTools)
                .defaultTools(classesTool, subjectTool, studentTool, gradeTool)
                //.defaultFunctions("getBookingDetails", "changeBooking", "cancelBooking") // FUNCTION CALLING
                .build();
    }

    @GetMapping(value = "/ai/generateStreamAsString", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> generateStreamAsString(String message) {
        Flux<String> content = chatClient.prompt()
                .system(s -> s.param("current_date", LocalDate.now().toString()))
                //.advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId).param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                .advisors(a -> a.param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                .user(message)
                //                .advisors(new QuestionAnswerAdvisor(vectorStore, SearchRequest.builder().query(message).build()))
                .stream()
                .content();
        return content
                .concatWith(Flux.just("[complete]"));
    }
//    @GetMapping(value = "/ai/generateStreamAsString", produces = MediaType.TEXT_PLAIN_VALUE)
//    public String generateStreamAsString(String message) {
//        String content = chatClient.prompt()
////                .system(s -> s.param("current_date", LocalDate.now().toString()))
////                .advisors(a -> a.param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
//                .user(message)
//                .call()
//                .content();
//        return content + "[complete]";
//    }
}
